nanopyx.methods.drift_alignment.estimator

  1import numpy as np
  2from math import sqrt
  3from scipy.interpolate import interp1d
  4
  5from .estimator_table import DriftEstimatorTable
  6from .corrector import DriftCorrector
  7from ...core.analysis.estimate_shift import GetMaxOptimizer
  8from ...core.utils.timeit import timeit
  9from ...core.analysis.ccm import calculate_ccm
 10from ...core.analysis.rcc import rcc
 11
 12
 13class DriftEstimator(object):
 14
 15    def __init__(self):
 16        self.estimator_table = DriftEstimatorTable()
 17        self.cross_correlation_map = None
 18        self.drift_xy = None
 19        self.drift_x = None
 20        self.drift_y = None
 21
 22    # @timeit
 23    def estimate(self, image_array, **kwargs):
 24        self.set_estimator_params(**kwargs)
 25
 26        n_slices = image_array.shape[0]
 27
 28        # x0, y0, x1, y1 correspond to the exact coordinates of the roi to be used or full image dims and should be a tuple
 29        if self.estimator_table.params["use_roi"] and self.estimator_table.params["roi"] is not None:  # crops image to roi
 30            print(self.estimator_table.params["use_roi"], self.estimator_table.params["roi"])
 31            x0, y0, x1, y1 = tuple(self.estimator_table.params["roi"])
 32            image_arr = image_array[:, y0:y1+1, x0:x1+1]
 33        else:
 34            image_arr = image_array
 35
 36        # checks time averaging, in case it's lower than 1 defaults to 1
 37        # if higher than n_slices/2 defaults to n_slices/2
 38        if self.estimator_table.params["time_averaging"] < 1:
 39            self.estimator_table.params["time_averaging"] = 1
 40        elif self.estimator_table.params["time_averaging"] > int(n_slices/2):
 41            self.estimator_table.params["time_averaging"] = int(n_slices/2)
 42
 43        # case of no temporal averaging
 44        if self.estimator_table.params["time_averaging"] == 1:
 45            image_averages = image_arr
 46        else: # case of temporal averaging
 47            # calculates number of time blocks for averaging
 48            image_averages = self.compute_temporal_averaging(image_arr)
 49
 50        method = self.estimator_table.params["shift_calc_method"]
 51
 52        if method == "rcc":
 53            shifts = rcc(image_averages, max_shift=self.estimator_table.params["max_expected_drift"])
 54            self.drift_x = shifts[0]
 55            self.drift_y = shifts[1]
 56        else:
 57            self.cross_correlation_map = np.array(calculate_ccm(np.array(image_averages).astype(np.float32), self.estimator_table.params["ref_option"]))
 58            max_shift = self.estimator_table.params["max_expected_drift"]
 59            if max_shift > 0 and max_shift*2+1 < self.cross_correlation_map.shape[1] and max_shift*2+1 < self.cross_correlation_map.shape[2]:
 60                ccm_x_start = int(self.cross_correlation_map.shape[1]/2 - max_shift)
 61                ccm_y_start = int(self.cross_correlation_map.shape[0]/2 - max_shift)
 62                slice_ccm = self.cross_correlation_map[ccm_y_start:ccm_y_start+(max_shift*2), ccm_x_start:ccm_x_start+(max_shift*2)]
 63            self.get_shifts_from_ccm()
 64
 65        if self.estimator_table.params["time_averaging"] > 1:
 66
 67            print("Interpolating time points")
 68            x_idx = np.linspace(1, image_array.shape[0], num=self.drift_x.shape[0], endpoint=True, dtype=int)
 69            x_interpolator = interp1d(x_idx, self.drift_x, kind="cubic") # linear seems to work similar as in nanoj-core however its codebase calls setInterpolation("Bicubic")
 70            self.drift_x = x_interpolator(range(1, image_array.shape[0] + 1))
 71            y_idx = np.linspace(1, image_array.shape[0], num=self.drift_y.shape[0], endpoint=True, dtype=int)
 72            y_interpolator = interp1d(y_idx, self.drift_y, kind="cubic") # linear seems to work similar as in nanoj-core however its codebase calls setInterpolation("Bicubic")
 73            self.drift_y = y_interpolator(range(1, image_array.shape[0] + 1))
 74
 75        self.drift_xy = []
 76        for i in range(image_array.shape[0]): 
 77            self.drift_xy.append(sqrt(pow(self.drift_x[i], 2) + pow(self.drift_y[i], 2)))
 78        self.drift_xy = np.array(self.drift_xy)
 79
 80        self.create_drift_table()
 81
 82        if self.estimator_table.params["apply"]:
 83            drift_corrector = DriftCorrector()
 84            drift_corrector.estimator_table = self.estimator_table
 85            tmp = drift_corrector.apply_correction(image_array)
 86            return tmp
 87        else:
 88            return None
 89
 90    def compute_temporal_averaging(self, image_arr):
 91        n_slices = image_arr.shape[0]
 92
 93        if self.estimator_table.params["use_roi"]:
 94            x0, y0, x1, y1 = self.estimator_table.params["roi"]
 95        else:
 96            x0, y0, x1, y1 = 0, 0, image_arr.shape[2]-1, image_arr.shape[1]-1
 97
 98        n_blocks = int(n_slices / self.estimator_table.params["time_averaging"])
 99        if (n_slices % self.estimator_table.params["time_averaging"]) != 0:
100            n_blocks += 1
101        image_averages = np.zeros((n_blocks, y1+1-y0, x1+1-x0))            
102        for i in range(n_blocks):
103            t_start = i * self.estimator_table.params["time_averaging"]
104            t_stop = (i + 1) * self.estimator_table.params["time_averaging"]
105            image_averages[i] = np.mean(image_arr[t_start:t_stop, :, :], axis=0)
106        return image_averages
107
108    def get_shift_from_ccm_slice(self, slice_index):
109        slice_ccm = self.cross_correlation_map[slice_index]
110
111        w = slice_ccm.shape[1]
112        h = slice_ccm.shape[0]
113
114        radius_x = w / 2.0
115        radius_y = h / 2.0
116
117        method = self.estimator_table.params["shift_calc_method"]
118
119        if method == "Max Fitting":
120            optimizer = GetMaxOptimizer(slice_ccm)
121            shift_y, shift_x = optimizer.get_max()
122        elif method == "Max":
123            shift_y, shift_x = np.unravel_index(slice_ccm.argmax(), slice_ccm.shape)
124
125        shift_x = round(radius_x - shift_x - 0.5, 3)
126        shift_y = round(radius_y - shift_y - 0.5, 3)
127
128        return (shift_x, shift_y)
129
130    def get_shifts_from_ccm(self):
131
132        drift_x = []
133        drift_y = []
134        drift = []
135
136        for i in range(self.cross_correlation_map.shape[0]):
137            drift.append(self.get_shift_from_ccm_slice(i))
138        
139        drift = np.array(drift)
140        drift_x = drift[:, 0]
141        drift_y = drift[:, 1]
142
143        bias_x = drift_x[0]
144        bias_y = drift_y[0]
145
146        self.drift_x = np.zeros((drift_x.shape[0]))
147        self.drift_y = np.zeros((drift_y.shape[0]))
148
149        for i in range(0, self.cross_correlation_map.shape[0]):
150            self.drift_x[i] = drift_x[i] - bias_x
151            self.drift_y[i] = drift_y[i] - bias_y
152            if self.estimator_table.params["ref_option"] == 1 and i > 0:
153                self.drift_x[i] += self.drift_x[i-1]
154                self.drift_y[i] += self.drift_y[i-1]
155
156        self.drift_x = np.array(self.drift_x)
157        self.drift_y = np.array(self.drift_y)
158
159    def create_drift_table(self):
160        table = []
161        for i in range(0, self.drift_xy.shape[0]):
162            table.append([self.drift_xy[i], self.drift_x[i], self.drift_y[i]])
163        table = np.array(table)
164        self.estimator_table.drift_table = table
165
166    def save_drift_table(self, save_as_npy=True, path=None):
167        if save_as_npy:
168            self.estimator_table.export_npy(path=path)
169        else:
170            self.estimator_table.export_csv(path=path)
171
172    def set_estimator_params(self, **kwargs):
173        self.estimator_table.set_params(**kwargs)